from langchain_ollama import OllamaLLM

from Knowledge.retriever import knowledge_search, qa_search
from Standard.query_knowledge import KnowledgeQuery
from Standard.metadata_query_helper import MetadataQueryHelper

# 初始化模型
llm = OllamaLLM(
    base_url="http://192.168.7.3:11434", model="qwen2.5:14b", temperature=0.2
)


def select_knowledge_base(query: str) -> str:
    route_prompt = f"""
    你是一个专业的知识库分类助手，任务是根据用户的问题内容，判断最适合的知识库进行回答（只能选择一个）。
    可选的知识库有：
    1. knowledge：概念解释类，例如污染种类、治理方法、计算规则等
    2. qa：标准问答类，包含日常常见问题的固定答复
    3. report：数据和图表类问题，例如销售、运营、统计报表等
    4. policy：政策标准类，例如制度、规则、规范、标准流程等

    请你只返回这四个名称中的一个（knowledge / qa / report / policy），不要输出任何多余内容。
    用户问题是：
    \"\"\"{query}\"\"\"
    """
    results = ""
    response = llm.invoke(route_prompt)
    result = response.strip().lower()
    if result not in {"knowledge", "qa", "report", "policy"}:
        result = "knowledge"
    if result == "knowledge":
        results = knowledge_search(query)
    elif result == "qa":
        results = qa_search(query)
    elif result == "policy":
        search = KnowledgeQuery()
        meta_extractor = MetadataQueryHelper()
        filters = meta_extractor.extract_metadata_filter(query)
        results = search.policy_search(query, filters=filters)
    return results
